package com.ailink.correctdistortiondemo

import android.graphics.Bitmap
import org.opencv.android.Utils
import org.opencv.core.Mat
import org.opencv.core.MatOfDouble
import org.opencv.imgproc.Imgproc

object DistortionCorrector {
    // 相机内参矩阵（需要根据实际相机标定结果修改）
    val cameraMatrix = arrayOf(
        doubleArrayOf(800.0, 0.0, 960.0),   // [fx, 0, cx]
        doubleArrayOf(0.0, 800.0, 540.0),   // [0, fy, cy]
        doubleArrayOf(0.0, 0.0, 1.0)

    )

    // 畸变系数（需要根据实际相机标定结果修改）
    val distCoeffs = doubleArrayOf(
        0.12,   // k1 径向畸变
        -0.25,  // k2
        0.001,  // p1 切向畸变
        0.003,  // p2
        0.15    // k3
    )

    fun correctDistortion(originalBitmap: Bitmap): Bitmap {
        // Step 1: 将Bitmap转换为OpenCV Mat格式
        val srcMat = Mat()
        Utils.bitmapToMat(originalBitmap, srcMat)

        // Step 2: 转换颜色空间（OpenCV默认使用BGR格式）
        Imgproc.cvtColor(srcMat, srcMat, Imgproc.COLOR_RGBA2BGR)

        // Step 3: 准备参数矩阵
        val cameraMatrixMat = Mat(3, 3, org.opencv.core.CvType.CV_64FC1)
        cameraMatrixMat.put(0, 0, *cameraMatrix[0])
        cameraMatrixMat.put(1, 0, *cameraMatrix[1])
        cameraMatrixMat.put(2, 0, *cameraMatrix[2])

        val distCoeffsMat = MatOfDouble(*distCoeffs)

        // Step 4: 执行畸变校正
        val dstMat = Mat()
        org.opencv.calib3d.Calib3d.undistort(
            srcMat,
            dstMat,
            cameraMatrixMat,
            distCoeffsMat
        )

        // Step 5: 转换回RGB格式
        Imgproc.cvtColor(dstMat, dstMat, Imgproc.COLOR_BGR2RGBA)

        // Step 6: 转换回Bitmap
        val resultBitmap = Bitmap.createBitmap(
            originalBitmap.width,
            originalBitmap.height,
            Bitmap.Config.ARGB_8888
        )
        Utils.matToBitmap(dstMat, resultBitmap)

        // 释放内存
        srcMat.release()
        cameraMatrixMat.release()
        distCoeffsMat.release()
        dstMat.release()

        return resultBitmap
    }
}